package main

import (
    "bytes"
    "encoding/json"
    "flag"
    "fmt"
    "log"
    "net/http"
    "strconv"
    "strings"
    "sync"
    "time"

    "golang.org/x/crypto/ssh"
)

var (
    apiToken string
)

type DeviceInfo struct {
    DeviceSN   string `json:"device_sn"`
    TargetUser string `json:"target_user"`
    TargetAddr string `json:"target_addr"`
    TargetPass string `json:"target_pass"`
}

type DispatchRequest struct {
    TaskType     string            `json:"task_type"`            // command|script
    TaskID       string            `json:"task_id"`              // 调用方自定义，必填
    Command      string            `json:"command,omitempty"`    // task_type=command 时必填
    ScriptURL    string            `json:"script_url,omitempty"` // task_type=script 时必填
    Mode         string            `json:"mode"`                 // sync|async|auto
    Timeout      int               `json:"timeout,omitempty"`    // 任务脚本自身超时
    SyncTimeout  int               `json:"sync_timeout,omitempty"`  // 同步窗口秒
    AsyncTimeout int               `json:"async_timeout,omitempty"` // 回调生命周期
    Interpreter  string            `json:"interpreter,omitempty"`
    Env          map[string]string `json:"env,omitempty"`
    CallbackURL  string            `json:"callback_url,omitempty"`
    JumpUser     string            `json:"jump_user"`
    JumpAddr     string            `json:"jump_addr"`
    JumpPass     string            `json:"jump_pass"`
    Devices      []DeviceInfo      `json:"devices"`
}

type DeviceResult struct {
    DeviceSN string `json:"device_sn"`
    TaskID   string `json:"task_id"`
    Status   string `json:"status"`            // running|done
    Stdout   string `json:"stdout,omitempty"`
    Stderr   string `json:"stderr,omitempty"`
    ExitCode *int   `json:"exit_code,omitempty"`
    Timeout  bool   `json:"timeout,omitempty"`
}

type RemoteResult struct {
    Status   string // "done" or "running"
    Stdout   string
    Stderr   string
    ExitCode *int
    Timeout  bool
}

func main() {
    listen := flag.String("listen", ":8080", "dispatch agent listen addr")
    flag.StringVar(&apiToken, "token", "", "Bearer authentication token (optional)")
    flag.Parse()

    http.HandleFunc("/dispatch", authMiddleware(handleDispatch))
    log.Printf("dispatch-agent listening on %s\n", *listen)
    if apiToken != "" {
        log.Printf("Authorization Bearer token is ENABLED")
    }
    log.Fatal(http.ListenAndServe(*listen, nil))
}

func authMiddleware(next http.HandlerFunc) http.HandlerFunc {
    return func(w http.ResponseWriter, r *http.Request) {
        if apiToken == "" {
            next(w, r)
            return
        }
        auth := r.Header.Get("Authorization")
        const prefix = "Bearer "
        if !strings.HasPrefix(auth, prefix) || strings.TrimSpace(auth[len(prefix):]) != apiToken {
            w.Header().Set("WWW-Authenticate", "Bearer")
            http.Error(w, "Unauthorized", http.StatusUnauthorized)
            return
        }
        next(w, r)
    }
}

func handleDispatch(w http.ResponseWriter, r *http.Request) {
    var req DispatchRequest
    if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
        http.Error(w, err.Error(), 400)
        return
    }
    if req.TaskType != "command" && req.TaskType != "script" {
        http.Error(w, "task_type must be command|script", 400)
        return
    }
    if req.TaskType == "command" && req.Command == "" {
        http.Error(w, "command empty", 400)
        return
    }
    if req.TaskType == "script" && req.ScriptURL == "" {
        http.Error(w, "script_url empty", 400)
        return
    }
    if (req.Mode == "async" || req.Mode == "auto") && req.CallbackURL == "" {
        http.Error(w, "async/auto need callback_url", 400)
        return
    }
    if req.TaskID == "" {
        http.Error(w, "task_id required", 400)
        return
    }

    results := make([]DeviceResult, len(req.Devices))
    var wg sync.WaitGroup

    for i, dev := range req.Devices {
        tid := req.TaskID // 全部节点同一 task_id
        results[i] = DeviceResult{
            DeviceSN: dev.DeviceSN,
            TaskID:   tid,
            Status:   "running",
        }
        wg.Add(1)
        devCopy := dev
        go func(i int, devInfo DeviceInfo, tid string) {
            defer wg.Done()
            res, err := runSSH(
                req.JumpUser, req.JumpAddr, req.JumpPass,
                devInfo.TargetUser, devInfo.TargetAddr, devInfo.TargetPass,
                req.TaskType, req.Command, req.ScriptURL,
                req.Interpreter, req.Env,
                req.Mode, req.Timeout, req.SyncTimeout, req.AsyncTimeout, req.CallbackURL,
                devInfo.DeviceSN, devInfo.TargetAddr, tid,
            )
            dr := DeviceResult{DeviceSN: devInfo.DeviceSN, TaskID: tid}
            if err != nil {
                dr.Status = "done"
                dr.Stderr = err.Error()
                code := -1
                dr.ExitCode = &code
            } else if res.Status == "done" {
                dr.Status = "done"
                dr.Stdout = res.Stdout
                dr.Stderr = res.Stderr
                dr.ExitCode = res.ExitCode
                dr.Timeout = res.Timeout
            } else {
                dr.Status = "running"
            }
            // 只在sync/auto下覆盖
            if req.Mode == "sync" || req.Mode == "auto" {
                results[i] = dr
            }
        }(i, devCopy, tid)
    }

    // sync/auto等所有goroutine完成（写入真实结果）
    isWait := req.Mode == "sync" || req.Mode == "auto"
    if isWait {
        wg.Wait()
    }
    // async直接返回初始running结构（已保证device_sn/task_id/status均有值）

    out := map[string]interface{}{"results": results}
    w.Header().Set("Content-Type", "application/json")
    json.NewEncoder(w).Encode(out)
}

func runSSH(
    jUser, jAddr, jPass,
    tUser, tAddr, tPass,
    taskType, command, scriptURL,
    interpreter string, env map[string]string,
    mode string, timeout, syncTimeout, asyncTimeout int, callbackURL, deviceSN, targetAddr, taskID string,
) (RemoteResult, error) {
    jcfg := &ssh.ClientConfig{
        User:            jUser,
        Auth:            []ssh.AuthMethod{ssh.Password(jPass)},
        HostKeyCallback: ssh.InsecureIgnoreHostKey(),
        Timeout:         10 * time.Second,
    }
    jcli, err := ssh.Dial("tcp", jAddr, jcfg)
    if err != nil {
        return RemoteResult{}, err
    }
    conn, err := jcli.Dial("tcp", tAddr)
    if err != nil {
        return RemoteResult{}, err
    }
    tcfg := &ssh.ClientConfig{
        User:            tUser,
        Auth:            []ssh.AuthMethod{ssh.Password(tPass)},
        HostKeyCallback: ssh.InsecureIgnoreHostKey(),
        Timeout:         10 * time.Second,
    }
    nc, chs, reqs, err := ssh.NewClientConn(conn, tAddr, tcfg)
    if err != nil {
        return RemoteResult{}, err
    }
    tcli := ssh.NewClient(nc, chs, reqs)
    defer tcli.Close()

    var evlist []string
    for k, v := range env {
        evlist = append(evlist, fmt.Sprintf("%s=%s", k, v))
    }
    evArg := strings.Join(evlist, ",")

    cb := callbackURL
    if cb != "" {
        sep := "?"
        if strings.Contains(cb, "?") {
            sep = "&"
        }
        cb = fmt.Sprintf("%s%stask_id=%s&device_sn=%s&target_addr=%s", callbackURL, sep, taskID, deviceSN, targetAddr)
    }

    var args []string
    args = append(args, "--sock", "/var/run/supervisor.sock")
    args = append(args, "--type", taskType)
    if taskType == "command" {
        args = append(args, "--command", command)
    } else {
        args = append(args, "--script-url", scriptURL)
    }
    args = append(args,
        "--interpreter", interpreter,
        "--env", evArg,
        "--mode", mode,
        "--timeout", fmt.Sprintf("%d", timeout),
        "--sync-timeout", fmt.Sprintf("%d", syncTimeout),
        "--async-timeout", fmt.Sprintf("%d", asyncTimeout),
        "--callback", cb,
        "--task-id", taskID,
    )
    cmdStr := "supervisor-client" + shellJoin(args...)

    sess, err := tcli.NewSession()
    if err != nil {
        return RemoteResult{}, err
    }
    defer sess.Close()

    var bout, berr bytes.Buffer
    sess.Stdout = &bout
    sess.Stderr = &berr

    runErr := sess.Run(cmdStr)
    if runErr != nil && mode == "sync" {
        return RemoteResult{}, runErr
    }

    var jr struct {
        TaskID   string `json:"task_id"`
        Stdout   string `json:"stdout"`
        Stderr   string `json:"stderr"`
        ExitCode *int   `json:"exit_code"`
        Timeout  bool   `json:"timeout"`
    }
    out := bout.Bytes()
    if err := json.Unmarshal(out, &jr); err == nil {
        if jr.ExitCode != nil {
            return RemoteResult{"done", jr.Stdout, jr.Stderr, jr.ExitCode, jr.Timeout}, nil
        }
        return RemoteResult{"running", "", "", nil, false}, nil
    }
    dummyCode := -1
    return RemoteResult{"done", "", berr.String(), &dummyCode, false}, nil
}

func shellJoin(args ...string) string {
    out := ""
    for _, a := range args {
        out += " " + strconv.Quote(a)
    }
    return out
}


